from src.module.TongyiModel import TongyiModel
from src.module.Embedding import Embedding

from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import Docx2txtLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.tools.retriever import create_retriever_tool
from langchain_community.tools.tavily_search import TavilySearchResults

from langchain import hub
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory

docx_path='C:/Users/AO237/Desktop/search/search_private.docx'
storage = {}

class AgentTools:
    def __init__(self):
        self.model = TongyiModel().model
        self.embedding = Embedding().embedding

    @staticmethod
    def get_session_history(session_id):
        if session_id not in storage:
            storage[session_id] = ChatMessageHistory()

        return storage[session_id]

    ## rag 检索本地的工具tool
    def rag_local_tool(self):
        # 拿到文件原始内容，进行切片
        loader = Docx2txtLoader(docx_path)
        docs = loader.load()
        documents = RecursiveCharacterTextSplitter(
            chunk_size=100,
            chunk_overlap=20
        ).split_documents(docs)

        # 拿到向量向量数据库句柄
        vector = FAISS.from_documents(documents, self.embedding)
        retriever = vector.as_retriever()

        # 创建两个工具，一个是和风的天气查询，一个是 rag 向量检索
        retriever_tools = create_retriever_tool(
            retriever,
            'local_search',
            '专门用于搜索公司内部文档。当问题涉及员工手册、产品文档、内部流程等明确记录在文档中的内容时，必须优先使用此工具。当在此工具中无法找到内容时，在尝试其它工具，如果都未找到内容，在自行根据实际场景回答。'
        )

        return retriever_tools

    def start(self):
        # 获取rag检索tool以及和风天气查询的 tool
        rag_tools = self.rag_local_tool()
        search = TavilySearchResults(max_results=1)

        # 创建 tools 列表
        tools = [search, rag_tools]

        # prompt 提示词模板
        prompt = hub.pull("hwchase17/openai-functions-agent")

        # 现在的模型还没办法直接调用工具，只是知道该调用哪个工具，此时。需要创建一个简单版  的 agent
        # agent 是一个智能体，可以将它想象为一个人，但是你要给它(model)、(tools)、(prompt)
        # agent 会自定执行 bind_tolls 方法，因此不需要在手动绑定 bind_tools
        agent = create_tool_calling_agent(self.model, tools, prompt)
        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False)

        # 现在就可以直接向它提问了, input 是因为 prompt 的模板定义的字段要求传递一个 input
        # hub.pull 时，不能使用本地代理，但是和风的接口是需要开启代理的
        # print(agent_executor.invoke({"input": "深圳的天气今天怎么样?"}))

        # 添加消息历史功能
        agent_with_message_history = RunnableWithMessageHistory(
            agent_executor,
            self.get_session_history,
            input_messages_key='input', ## 需要固定，表示问题的字段名称，不可更换
            history_messages_key='chat_history' # 需要固定，此字段表示历史字段的占位符
        )

        response = agent_with_message_history.invoke({
            'input': '什么是 ”aafe123?“ '
        }, config={"configurable": {'session_id': "123"}})

        print(response)



